#include "aodv_forward.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "aodv_rreq.h"
#include "aodv_rerr.h"  // 添加RERR处理头文件
#include "aodv_socket.h"
#include "aodv_packet.h"  // 添加数据包处理头文件
#include <string.h>

static forward_stats_t *g_stats = NULL;

// 初始化数据转发模块
void forward_init(forward_stats_t *stats) {
    g_stats = stats;
    if (g_stats) {
        memset(g_stats, 0, sizeof(forward_stats_t));
    }
}

// 生成数据包ID
uint32_t generate_packet_id(void) {
    static uint32_t packet_counter = 0;
    return ++packet_counter;
}

// 检查是否需要路由发现
int needs_route_discovery(const aodv_routing_table *routing_table, uint32_t dest_addr) {
    if (routing_table == NULL) return 1;
    
    // 检查是否有到目标的有效路由
    routing_table_entry *route = routing_table_find((aodv_routing_table*)routing_table, dest_addr);
    if (route == NULL) {
        printf("No route to %s, need route discovery\n", 
               inet_ntoa(*(struct in_addr*)&dest_addr));
        return 1; // 没有路由，需要发现
    }
    
    if (!routing_table_is_valid((aodv_routing_table*)routing_table, dest_addr)) {
        printf("Route to %s is invalid, need route discovery\n", 
               inet_ntoa(*(struct in_addr*)&dest_addr));
        return 1; // 路由无效，需要重新发现
    }
    
    printf("Valid route to %s exists, no need for route discovery\n", 
           inet_ntoa(*(struct in_addr*)&dest_addr));
    return 0; // 有有效路由，不需要发现
}

// 执行路由发现
forward_result_t perform_route_discovery(int socket, uint32_t src_addr, uint32_t dest_addr,
                                       aodv_routing_table *routing_table,
                                       uint32_t *sequence_number, uint16_t *rreq_id) {
    if (g_stats) g_stats->route_discovery_requests++;
    
    printf("Starting route discovery to %s\n", 
           inet_ntoa(*(struct in_addr*)&dest_addr));
    
    // 生成并发送RREQ
    aodv_rreq rreq;
    uint32_t rreq_dest_seq = 0; // 未知目标序列号
    
    // 如果有序列号，使用它
    routing_table_entry *dest_route = routing_table_find(routing_table, dest_addr);
    if (dest_route) {
        rreq_dest_seq = dest_route->dest_seq_num;
    }
    
    build_rreq_packet(&rreq, *rreq_id, dest_addr, rreq_dest_seq, src_addr, (*sequence_number)++, 0);
    
    // 广播RREQ
    int result = aodv_socket_send(socket, &rreq, sizeof(aodv_rreq), 
                                  inet_addr("255.255.255.255"), AODV_PORT);
    
    if (result <= 0) {
        fprintf(stderr, "Failed to send RREQ for route discovery\n");
        if (g_stats) g_stats->route_discovery_failures++;
        return FORWARD_SEND_FAILED;
    }
    
    printf("RREQ sent for route discovery to %s\n", 
           inet_ntoa(*(struct in_addr*)&dest_addr));
    return FORWARD_ROUTE_DISCOVERY_NEEDED;
}

// 转发数据包到下一跳
forward_result_t forward_packet_to_next_hop(int socket, const data_packet_header_t *header,
                                          const void *data, size_t data_len,
                                          const aodv_routing_table *routing_table) {
    if (header == NULL || data == NULL || routing_table == NULL) {
        return FORWARD_INVALID_PACKET;
    }
    
    // 查找下一跳
    routing_table_entry *route = routing_table_find((aodv_routing_table*)routing_table, header->dest_addr);
    if (route == NULL || !routing_table_is_valid((aodv_routing_table*)routing_table, header->dest_addr)) {
        printf("No valid route to %s\n", 
               inet_ntoa(*(struct in_addr*)&header->dest_addr));
        return FORWARD_NO_ROUTE;
    }
    
    // 构建完整的数据包（头部 + 数据）
    size_t total_len = sizeof(data_packet_header_t) + data_len;
    uint8_t *packet = malloc(total_len);
    if (packet == NULL) {
        fprintf(stderr, "Failed to allocate memory for packet\n");
        return FORWARD_SEND_FAILED;
    }
    
    // 复制头部和数据
    memcpy(packet, header, sizeof(data_packet_header_t));
    memcpy(packet + sizeof(data_packet_header_t), data, data_len);
    
    // 发送到下一跳
    int result = aodv_socket_send(socket, packet, total_len, 
                                  route->next_hop_addr, AODV_PORT + 1); // 数据端口使用AODV端口+1
    
    free(packet);
    
    if (result <= 0) {
        fprintf(stderr, "Failed to forward packet to next hop %s, triggering RERR\n", 
                inet_ntoa(*(struct in_addr*)&route->next_hop_addr));
        
        // 发送失败，触发RERR消息
        handle_route_error(route->next_hop_addr, socket, (aodv_routing_table*)routing_table, 
                          ((aodv_routing_table*)routing_table)->entries[0].dest_addr); // 使用任意本地IP
        
        return FORWARD_SEND_FAILED;
    }
    
    printf("Successfully forwarded packet to %s via next hop %s\n", 
           inet_ntoa(*(struct in_addr*)&header->dest_addr),
           inet_ntoa(*(struct in_addr*)&route->next_hop_addr));
    
    if (g_stats) g_stats->packets_forwarded++;
    return FORWARD_SUCCESS;
}

// 发送数据包（集成路由发现）
forward_result_t forward_data_packet(int socket, uint32_t src_addr, uint32_t dest_addr,
                                   const void *data, size_t data_len,
                                   aodv_routing_table *routing_table,
                                   uint32_t *sequence_number, uint16_t *rreq_id) {
    if (data == NULL || routing_table == NULL || sequence_number == NULL || rreq_id == NULL) {
        return FORWARD_INVALID_PACKET;
    }
    
    char src_str[INET_ADDRSTRLEN];
    char dest_str[INET_ADDRSTRLEN];
    strcpy(src_str, inet_ntoa(*(struct in_addr*)&src_addr));
    strcpy(dest_str, inet_ntoa(*(struct in_addr*)&dest_addr));
    printf("Attempting to send data packet from %s to %s, size=%zu\n",
           src_str, dest_str, data_len);
    
    // 检查是否需要路由发现
    if (needs_route_discovery(routing_table, dest_addr)) {
        printf("No valid route available, initiating route discovery\n");
        return perform_route_discovery(socket, src_addr, dest_addr, routing_table, 
                                      sequence_number, rreq_id);
    }
    
    // 有有效路由，直接发送数据包
    data_packet_header_t header;
    header.src_addr = src_addr;
    header.dest_addr = dest_addr;
    header.packet_id = generate_packet_id();
    header.packet_type = PACKET_TYPE_DATA;
    header.data_len = data_len;
    header.timestamp = (uint32_t)time(NULL);
    
    forward_result_t result = forward_packet_to_next_hop(socket, &header, data, data_len, routing_table);
    
    if (result == FORWARD_SUCCESS) {
        if (g_stats) g_stats->packets_sent++;
        printf("Data packet sent successfully (ID=%u)\n", header.packet_id);
    }
    
    return result;
}

// 处理接收到的数据包
forward_result_t process_received_packet(int socket, const void *packet, size_t packet_len,
                                       uint32_t source_ip, aodv_routing_table *routing_table,
                                       uint32_t local_ip, uint32_t *sequence_number, uint16_t *rreq_id) {
    if (packet == NULL || routing_table == NULL || packet_len < sizeof(data_packet_header_t)) {
        return FORWARD_INVALID_PACKET;
    }
    
    // 解析数据包头
    const data_packet_header_t *header = (const data_packet_header_t *)packet;
    const void *data = (const uint8_t *)packet + sizeof(data_packet_header_t);
    size_t data_len = packet_len - sizeof(data_packet_header_t);
    
    printf("Received packet from %s to %s, ID=%u, size=%zu\n",
           inet_ntoa(*(struct in_addr*)&header->src_addr),
           inet_ntoa(*(struct in_addr*)&header->dest_addr),
           header->packet_id, data_len);
    
    // 检查是否是发给自己的数据包
    if (header->dest_addr == local_ip) {
        printf("Packet arrived at destination! ID=%u, Data: ", header->packet_id);
        // 打印数据内容（假设是文本）
        const char *text_data = (const char *)data;
        for (size_t i = 0; i < data_len && i < 50; i++) { // 最多打印50个字符
            if (text_data[i] >= 32 && text_data[i] <= 126) {
                printf("%c", text_data[i]);
            } else {
                printf(".");
            }
        }
        printf("\n");
        
        if (g_stats) g_stats->packets_received++;
        return FORWARD_SUCCESS;
    }
    
    // 不是发给自己的，需要转发
    printf("Packet not for this node, forwarding...\n");
    
    // 检查是否有到目标的路由
    if (needs_route_discovery(routing_table, header->dest_addr)) {
        printf("No route to destination, cannot forward\n");
        return FORWARD_NO_ROUTE;
    }
    
    // 转发数据包
    return forward_packet_to_next_hop(socket, header, data, data_len, routing_table);
}

// 获取转发统计
void get_forward_stats(const forward_stats_t *stats, char *buffer, size_t buffer_size) {
    if (stats == NULL || buffer == NULL) return;
    
    snprintf(buffer, buffer_size,
              "=== Data Forwarding Statistics ===\n"
              "Packets Sent: %u\n"
              "Packets Received: %u\n"
              "Packets Forwarded: %u\n"
              "Route Discovery Requests: %u\n"
              "Route Discovery Failures: %u\n"
              "Success Rate: %.1f%%\n",
              stats->packets_sent,
              stats->packets_received,
              stats->packets_forwarded,
              stats->route_discovery_requests,
              stats->route_discovery_failures,
              stats->packets_sent > 0 ? 
              (float)(stats->packets_sent - stats->route_discovery_failures) * 100.0f / stats->packets_sent : 0.0f);
}

// 清理转发模块
void forward_cleanup(void) {
    g_stats = NULL;
}